/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#include "tee_teleport.h"
#include <sys/stat.h>
#include <unistd.h>
#include <ctype.h>
#include <fcntl.h>
#include <limits.h>
#include <stdint.h>

#include "dir.h"

enum CliArgType {
    ARG_CLEAN,
    ARG_GRPID,
    ARG_NSID_SET,
    ARG_CONF_RES,
    ARG_VMID,
    ARG_MEMORY,
    ARG_DISK_SIZE,
    ARG_CPUSET,
    ARG_CPUS,
    ARG_CONTAINER,
    ARG_CONF_CONT
};

static const struct option g_toolOptions[] = {
                                              {"help",              required_argument, NULL, 'h'},
                                              {"clean",             no_argument,       NULL,  ARG_CLEAN},
                                              {"grpid",             required_argument, NULL,  ARG_GRPID},
                                              {"nsid",              required_argument, NULL,  ARG_NSID_SET},
                                              {"config-container",  no_argument,       NULL,  ARG_CONF_CONT},
                                              {"containerid",       required_argument, NULL,  ARG_CONTAINER},
                                              {"config-resource",   no_argument,       NULL,  ARG_CONF_RES},
                                              {"memory",            required_argument, NULL,  ARG_MEMORY},
                                              {"cpuset-cpus",       required_argument, NULL,  ARG_CPUSET},
                                              {"cpus",              required_argument, NULL,  ARG_CPUS},
                                              {"disk-size",         required_argument, NULL,  ARG_DISK_SIZE},
                                              {"vmid",              required_argument, NULL,  ARG_VMID},
                                              {NULL, 0, NULL, 0}};

static int32_t PrintUsage(const struct TeeTeleportArgs *args, uint32_t sessionID)
{
    (void)args;
    (void)sessionID;
    printf("This is a tool for running elf, java or python apps in iTrustee.\n"
           "Usage:\n"
           "--memory: set memory size for iTrustee corresponds to ree container\n"
           "--disk-size: set tmpfs size for iTrustee corresponds to ree container\n"
           "--cpuset-cpus: set cpuset for iTrustee corresponds to ree container\n"
           "--cpus: set cpus usage limit for iTrustee corresponds to ree container\n"
           "--nsid: enter nsid and perform subsequent operations.\n"
           "--config-resource: specify nsid for iTrustee corresponds to ree container\n"
           "--config-container: writes information to a triplet on the tee side through nsid and containerid.\n"
           "--containerid: set container id for other actions\n"
           "--clean: clean entered nsid and containerid\n"
           "--vmid: set vmid for iTrustee corresponds to ree vm\n"
           "--grpid: pass grpid for iTrustee corresponds to ree container\n"
           "-h: print this help message\n");
    return 0;
}

static int32_t ParseStr(const char *str, int base, long *num, char **endPtr)
{
    errno = 0;
    long res = strtol(str, endPtr, base);
    if ((errno == ERANGE && (res == LONG_MAX || res == LONG_MIN))
            || (errno != 0 && res == 0) || *endPtr == str) {
        printf("failed to parse str %s\n", str);
        return -EINVAL;
    }

    *num = res;
    return 0;
}

static int32_t IsMemConfigValid(const char *memSize)
{
    if (memSize == NULL || strlen(memSize) == 0)
        return -EINVAL;
    
    char *endPtr = NULL;
    long num;
    int32_t ret = ParseStr(memSize, 10, &num, &endPtr);
    if (ret != 0) {
        printf("failed to parse mem config, mem size is %s\n", memSize);
        return ret;
    }

    if (endPtr == NULL) {
        printf("please add k/m/g after mem size!, memSize is %s\n", memSize);
        return -EINVAL;
    }

    if (*endPtr != 'K' && *endPtr != 'k' &&
        *endPtr != 'M' && *endPtr != 'm' &&
        *endPtr != 'G' && *endPtr != 'g') {
        printf("please add k/m/g after mem size!\n");
        return -EINVAL;
    }

    return 0;
}

/* containerMsg is a combination of containerId and namespace_id */
static int32_t CheckContainerId(const char *containerMsg)
{
    size_t idLen = strlen(containerMsg);
    /* check length */
    if (idLen != CONTAINER_ID_LEN) {
        printf("idLen = %d\n", (int)idLen);
        return -EINVAL;
    }

    for (size_t i = 0; i < CONTAINER_ID_LEN; i++) {
        if (!isxdigit(containerMsg[i])) {
            printf("not hex parameter\n");
            return -EINVAL;
        }
    }
    return 0;
}

static int32_t DoCleanContainerMsg(const struct TeeTeleportArgs *args, uint32_t sessionID)
{
    (void)sessionID;
    errno_t ret;
    struct TeePortalContainerType config = {0};

    if (CheckContainerId(args->containerMsg) != 0) {
        printf("bad container mesg!\n");
        return -EINVAL;
    }
    ret = memcpy_s(config.containerid, sizeof(config.containerid), args->containerMsg, CONTAINER_ID_LEN);
    if (ret != 0) {
        printf("memcpy_s failed, ret = %u", ret);
        return -EINVAL;
    }

    return TeeSendContainerMsg(&config, CONTAINER_STOP);
}

static int32_t DoCleanGroupId(const struct TeeTeleportArgs *args, uint32_t sessionID)
{
    (void)sessionID;
    int32_t ret = -1;
    printf("grp id is %s\n", args->grpId);

    long grpId = -1;
    char *endPtr = NULL;
    int base = 10;
    ret = ParseStr(args->grpId, base, &grpId, &endPtr);
    if (ret != 0 || *endPtr != '\0') {
        printf("failed to parse group id when clean!\n");
        return ret;
    }

    ret = TeeClean((int)grpId);
    if (ret != 0) {
        printf("failed to clean tee resource!\n");
    }

    return ret;
}

static int32_t DoClean(const struct TeeTeleportArgs *args, uint32_t sessionID)
{
    int32_t ret = 0;
    if (strlen(args->grpId) != 0) {
        ret = DoCleanGroupId(args, sessionID);
        if (ret != 0) {
            printf("failed to clean group id!\n");
            return ret;
        }
    }

    if (strlen(args->containerMsg) != 0) {
        ret = DoCleanContainerMsg(args, sessionID);
        if (ret != 0) {
            printf("failed to clean container msg!\n");
        }
    }
    return ret;
}

static int32_t DoRconfig(const struct TeeTeleportArgs *args, uint32_t sessionID)
{
    (void)sessionID;
    int32_t ret = -1;
    
    struct TeePortalRConfigType config = {0};
    char *endPtr = NULL;
    int base = 10;
    long nsid = 0;
    ret = ParseStr(args->nsId, base, &nsid, &endPtr);
    if (ret == -EINVAL || *endPtr != '\0') {
        printf("failed to parse nsid!, nsid is %ld\n", nsid);
        return ret;
    }

    endPtr = NULL;
    if (strcmp(args->cpus, "") != 0 && strcpy_s(config.cpus, PARAM_LEN, args->cpus) == 0) {
        printf("succeed to copy cpus config!\n");
    }

    if (IsMemConfigValid(args->mem) == 0 && strcpy_s(config.memSize, PARAM_LEN, args->mem) == 0) {
        printf("succeed to parse mem size config!\n");
    }

    if (IsMemConfigValid(args->diskSize) == 0 && strcpy_s(config.diskSize, PARAM_LEN, args->diskSize) == 0) {
        printf("succeed to parse disk size config!\n");
    }

    if (strcmp(args->cpuset, "") != 0 && strcpy_s(config.cpuset, PARAM_LEN, args->cpuset) == 0) {
        printf("succeed to parse cpuset config!\n");
    }
    
    ret = TeeRconfig(&config, nsid);
    printf("TeeRconfig ret is %d\n", ret);
    return ret;
}

static int32_t CheckParams(const struct TeeTeleportArgs *args, struct TeePortalContainerType *config)
{
    int32_t ret;
    char *endPtr = NULL;
    int base = 10; // 10 is the normal length of nsid
    long ns_id = 0;

    ret = ParseStr(args->nsId, base, &ns_id, &endPtr);
    if (ret == -EINVAL || *endPtr != '\0') {
        printf("failed to parse nsid! nsid is %d\n", config->nsid);
        return ret;
    }

    config->nsid = (uint32_t)ns_id;
    if (CheckContainerId(args->containerMsg) != 0) {
        printf("bad container mesg!\n");
        return -EINVAL;
    }

    ret = memcpy_s(config->containerid, sizeof(config->containerid), args->containerMsg, CONTAINER_ID_LEN);
    if (ret != 0) {
        printf("memcpy_s failed, ret = %u", ret);
        return -EINVAL;
    }

    return 0;
}

static int32_t DoContainer(const struct TeeTeleportArgs *args, uint32_t sessionID)
{
    (void)sessionID;
    struct TeePortalContainerType config = { 0 };
    if (CheckParams(args, &config) != 0) {
        printf("bad container mesg!\n");
        return -EINVAL;
    }

    return TeeSendContainerMsg(&config, CONTAINER_OPEN);
}

static const struct TeeTeleportFunc g_teleportFuncTable[] = {
    {TP_HELP,       PrintUsage,  false},
    {TP_CONF_CONT,  DoContainer, false},
    {TP_CLEAN,      DoClean,     false},
    {TP_CONF_RES,   DoRconfig,   false},
};

static const uint32_t g_teleportFuncNum = sizeof(g_teleportFuncTable) / sizeof(g_teleportFuncTable[0]);

static int32_t HandleUserCmd(struct TeeTeleportArgs *args)
{
    uint32_t sessionID = 0;
    for (uint32_t i = 0; i < g_teleportFuncNum; i++) {
        if (args->cmd[g_teleportFuncTable[i].type]) {
            return g_teleportFuncTable[i].func(args, sessionID);
        }
    }
    printf("cannot handle user cmd\n");
    return -EINVAL;
}

static int32_t CopyArgs(bool *argkey, char *argvalue, uint32_t size)
{
    *argkey = true;
    if (argvalue == NULL)
        return 0;

    size_t optLen = strnlen(optarg, PATH_MAX);
    if (optLen == 0 || optLen >= PATH_MAX) {
        printf("opt arg is invalid\n");
        return -EINVAL;
	}

    if (strcpy_s(argvalue, size - 1, optarg) != EOK) {
        printf("strcpy_s failed\n");
        return -EFAULT;
    }
    return 0;
}

#define CASE_COPY_ARG(opt, argkey, argvalue, size) \
    case opt: \
        if (CopyArgs(argkey, argvalue, size) != 0) \
            return -EFAULT; \
        break;

static int32_t ParseArgs(int32_t argc, char **argv, struct TeeTeleportArgs *args)
{
    while (1) {
        int32_t optIndex = 0;
        int32_t opt = getopt_long(argc, argv, "h", g_toolOptions, &optIndex);
        if (opt < 0)
            break;
        switch (opt) {
        CASE_COPY_ARG('h',             &args->cmd[TP_HELP],      NULL,                  0);
        CASE_COPY_ARG(ARG_CLEAN,       &args->cmd[TP_CLEAN],     NULL,                  0);
        CASE_COPY_ARG(ARG_GRPID,       &args->cmd[TP_GRPID],     args->grpId,           PARAM_LEN);
        CASE_COPY_ARG(ARG_NSID_SET,    &args->cmd[TP_NSID_SET],  args->nsId,            PARAM_LEN);
        CASE_COPY_ARG(ARG_CONF_RES,    &args->cmd[TP_CONF_RES],  NULL,                  0);
        CASE_COPY_ARG(ARG_CONTAINER,   &args->cmd[TP_CONTAINER], args->containerMsg,    PARAM_LEN_MAX);
        CASE_COPY_ARG(ARG_CONF_CONT,   &args->cmd[TP_CONF_CONT], NULL,                  PARAM_LEN_MAX);
        CASE_COPY_ARG(ARG_MEMORY,      &args->cmd[TP_MEM],       args->mem,             PARAM_LEN);
        CASE_COPY_ARG(ARG_CPUSET,      &args->cmd[TP_CPUSET],    args->cpuset,          PARAM_LEN_MAX);
        CASE_COPY_ARG(ARG_CPUS,        &args->cmd[TP_CPUS],      args->cpus,            PARAM_LEN);
        CASE_COPY_ARG(ARG_DISK_SIZE,   &args->cmd[TP_DISK_SIZE], args->diskSize,        PARAM_LEN);
        default:
            printf("please use -h to see the usage.\n");
            return -EINVAL;
        }
    }
    return 0;
}

#define MIN_ARG_COUNT 2
int32_t main(int32_t argc, char *argv[])
{
    struct TeeTeleportArgs args = { 0 };
    int32_t ret = 0;

    if (argc < MIN_ARG_COUNT) {
        printf("please use \"tee_teleport -h\" for more help!\n");
        return ret;
    }

    if (ParseArgs(argc, argv, &args) != 0)
        return -EINVAL;

    if (InitPortal() != 0) {
        printf("cannot init portal\n");
        return -EFAULT;
    }

    ret = HandleUserCmd(&args);
    DestroyPortal();
    return ret;
}

